import os
import cv2

# =============================================
#               数据集图像扩充
# =============================================

path = './JPEGImages/'
savepath = './new/'
imagewight = 640	# 这两个参数改为自己需要的长和宽
imagehight = 480
imagerow = 30
print("image read from {}, save in {}".format(path, savepath))

path_list = os.listdir(path)

imagenumber = 0
for filename in path_list:
    print(filename)
    image = cv2.imread(os.path.join(path,filename), 1)

    # =============================================
    #               固定尺寸缩放
    # =============================================
    res = cv2.resize(image, (imagewight, imagehight), interpolation=cv2.INTER_LINEAR)
    imagename = str(imagenumber)+'.jpg'
    cv2.imwrite(os.path.join(savepath,imagename), res)
    imagenumber += 1

    # =============================================
    #              上下翻转
    # =============================================
    dst1 = cv2.flip(res, 0)
    imagename = str(imagenumber) + '.jpg'
    cv2.imwrite(os.path.join(savepath, imagename), dst1)
    imagenumber += 1

    # =============================================
    #               左右对称翻转
    # =============================================
    dst3 = cv2.flip(res, -1)
    imagename = str(imagenumber) + '.jpg'
    cv2.imwrite(os.path.join(savepath, imagename), dst3)
    imagenumber += 1

    # =============================================
    #               仿射旋转
    # =============================================
    # 这里的第一个参数为旋转中心，第二个为旋转角度，第三个为旋转后的缩放因子
    # 可以通过设置旋转中心，缩放因子以及窗口大小来防止旋转后超出边界的问题
    M = cv2.getRotationMatrix2D((imagewight/2 , imagehight/2), imagerow, 1.2)
    wap = cv2.warpAffine(res, M, (imagewight, imagehight), borderValue=(255, 255, 255))  # M为上面的旋转矩阵
    imagename = str(imagenumber) + '.jpg'
    cv2.imwrite(os.path.join(savepath, imagename), wap)
    if imagerow < 320:
        imagerow += 30
    else:
        imagerow = 30
    imagenumber += 1

print("all image secussful!! image save in {}".format(savepath))
cv2.waitKey()
cv2.destroyAllWindows()
